import torch
from torch.utils.data import DataLoader
from torchnet import meter
from torchvision import datasets, transforms
from tqdm import tqdm

from model.GooleNet import GooLeNet
from utils.config import opt

transform = transforms.Compose([transforms.Resize((224,224)),transforms.Grayscale(num_output_channels=1),
                                transforms.ToTensor(),
                                transforms.Normalize(mean=[0.2458], std=[0.0612])])
test_data = datasets.ImageFolder('./data/test',transform=transform)
test_loader = DataLoader(dataset=test_data,batch_size=8,shuffle=False)

device = ("cuda:0" if torch.cuda.is_available() else "cpu")

def test(dataloader):
    model = GooLeNet(num_classes=2,aux_logits=True)
    if opt.load_model_path:
        model.load(opt.load_model_path)
    model.to(device)
    confusion_matrix = meter.ConfusionMeter(2)
    for ii, (val_input, label) in enumerate(tqdm(dataloader)):
        val_input = val_input.to(device)
        label = label.to(device)
        _, prediction = torch.max(model(val_input).data, dim=1)
        confusion_matrix.add(prediction.detach(), label.detach())

    model.train()
    cm_value = confusion_matrix.value()
    accuracy = 100. * (cm_value[0][0] + cm_value[1][1]) / (cm_value.sum())
    print(cm_value)
    print(accuracy)


if __name__ == '__main__':
    test(test_loader)